'''
Brief:
  CNN模型代码
Requirement:
  完善ClassifierNet类，要求实现对于数字1-5的分类，导出模型可以的类输出应该为置信度(0-1)
'''
import torch
import torch.nn as nn
import torch.nn.functional as F
class ClassifierNet(nn.Module):
    '''
    Brief:
      输入为52x52的图像，由于pytorch必须要把channel放到第二个维度，因此需要变为1
      x52x52
    '''
    def __init__(self,num_classes = 8,use_bp=True):
        """
        Args:
          num_classes(int):分类数目，默认值为8
          use_bp(bool):是否使用反向传播，这个与训练和导出模型有关。
            这里主要是因为torch的softmax层输出对应的损失函数我忘了，如果直接输出用交叉熵函数即可
        """
        super(ClassifierNet,self).__init__()
        self.seq_module = nn.Sequential()
        self.num_classes = num_classes
        self.use_bp = use_bp
        #init model
        self.seq_module.add_module(
                "conv_1",
                nn.Conv2d(in_channels=1,out_channels=8,kernel_size=5,stride=1,padding=2,bias=True))
        #same padding,output = [-1,16,28,28]
        self.seq_module.add_module("leaky_1", nn.LeakyReLU(0.1))
        self.seq_module.add_module('max_pool_1',nn.MaxPool2d(2,2))
        #downsample [-1,16,14,14]
        #output = [(input_size - kernel_size + 2*padding) / stride] + 1
        self.seq_module.add_module("conv_2",
                nn.Conv2d(in_channels=8,out_channels=16,kernel_size=5,stride=1,padding=2,bias=False))
        self.seq_module.add_module('bn_2',nn.BatchNorm2d(16,momentum=0.9, eps=1e-5))
        self.seq_module.add_module('leaky_2',nn.LeakyReLU(0.1))
        self.seq_module.add_module('max_pool_2',nn.MaxPool2d(2,2))
        #tensor [-1,32,7,7]
        #self.fc1 = nn.Linear(16*13*13, 26)
        self.fc2 = nn.Linear(16*13*13,self.num_classes)
        self.softmax = nn.Softmax(dim=1)
    
    def forward(self,x):
        batch_size = x.size(0)
        output = self.seq_module(x)

        output = output.view([batch_size,-1])#你妈的，必须赋值
        #output = self.fc1(output)
        #output = F.relu(output)

        output = self.fc2(output)
        if not self.use_bp:
			#如果不使用反向传播
        	output = self.softmax(output)
        #直接输出，计算loss直接用CrossEntrypy，同时计算softmax+log+NLLloss
        #如果使用Logsoftmax，那么直接用NLLloss
        return output


if __name__ == '__main__':
    """
    导出libtorch模型
    """
    model = ClassifierNet(use_bp=False)
    x = torch.ones((1,1,52,52))
    checkpoint = torch.load("")
    model.load_state_dict(checkpoint['state_dict'])
    trace_model = torch.jit.trace(model,x)
    trace_model.save("cnn.pt")

    #导出onnx模型：
    #torch.onnx.export(model,x,
    #                  "cnn.onnx",
    #                  #export_params=True,
    #                  opset_version=10,
    #                  do_constant_folding=True,
    #                  input_names=['input'], output_names=['output'])
    

